from fastapi import FastAPI, File, UploadFile, HTTPException, Header, Query
from fastapi.responses import StreamingResponse
import httpx
import logging
from typing import List, Optional
import os
from pydantic import BaseModel

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s.%(msecs)03d [%(levelname)s] %(name)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    filename='reinforce_proxy.log',
    filemode='a',
    #encoding='utf-8'
)
logger = logging.getLogger('reinforce_proxy')

# 配置FastAPI的访问日志
access_logger = logging.getLogger("uvicorn.access")
access_logger.setLevel(logging.INFO)

app = FastAPI(title="APK加固代理服务")

# 目标服务器基础URL
TARGET_BASE_URL = "https://isec.iflytek.com"

# 请求模型
class AccessTokenRequest(BaseModel):
    clientKey: str
    expire: int

class TaskStatusRequest(BaseModel):
    __root__: List[int]

class ReinforceTaskRequest(BaseModel):
    dataSrc: str
    fileName: str
    downloadName: str
    fileSaveName: str
    filePath: str
    minioFileUrl: str
    md5: str
    userAccount: str
    keystoreId: int

@app.post("/sec/open/auth/accessToken")
async def get_access_token(request: AccessTokenRequest):
    """获取访问令牌"""
    async with httpx.AsyncClient() as client:
        try:
            response = await client.post(
                f"{TARGET_BASE_URL}/sec/open/auth/accessToken",
                json=request.dict()
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"获取访问令牌失败: {str(e)}")
            raise HTTPException(status_code=500, detail=str(e))

@app.post("/sec/open/file/upload")
async def upload_file(
    file: UploadFile = File(...),
    authorization: str = Header(...)
):
    """上传APK文件"""
    async with httpx.AsyncClient() as client:
        try:
            files = {"file": (file.filename, file.file, file.content_type)}
            headers = {"Authorization": authorization}
            response = await client.post(
                f"{TARGET_BASE_URL}/sec/open/file/upload",
                files=files,
                headers=headers
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"上传文件失败: {str(e)}")
            raise HTTPException(status_code=500, detail=str(e))

@app.post("/sec/open/reinforceCommon/addTask")
async def create_reinforce_task(
    request: ReinforceTaskRequest,
    authorization: str = Header(...)
):
    """创建加固任务"""
    async with httpx.AsyncClient() as client:
        try:
            headers = {
                "Authorization": authorization,
                "Content-Type": "application/json;charset=UTF-8"
            }
            response = await client.post(
                f"{TARGET_BASE_URL}/sec/open/reinforceCommon/addTask",
                json=request.dict(),
                headers=headers
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"创建加固任务失败: {str(e)}")
            raise HTTPException(status_code=500, detail=str(e))

@app.post("/sec/open/reinforceCommon/taskStatus")
async def check_task_status(
    request: TaskStatusRequest,
    authorization: str = Header(...)
):
    """查询任务状态"""
    async with httpx.AsyncClient() as client:
        try:
            headers = {
                "Authorization": authorization,
                "Content-Type": "application/json;charset=UTF-8"
            }
            response = await client.post(
                f"{TARGET_BASE_URL}/sec/open/reinforceCommon/taskStatus",
                json=request.__root__,
                headers=headers
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"查询任务状态失败: {str(e)}")
            raise HTTPException(status_code=500, detail=str(e))

@app.get("/sec/open/reinforceCommon/download")
async def download_reinforced_apk(
    taskId: int = Query(...),
    authorization: str = Header(...)
):
    """下载加固后的APK"""
    async with httpx.AsyncClient() as client:
        try:
            headers = {"Authorization": authorization}
            response = await client.get(
                f"{TARGET_BASE_URL}/sec/open/reinforceCommon/download",
                params={"taskId": taskId},
                headers=headers
            )
            response.raise_for_status()
            
            # 检查响应类型
            content_type = response.headers.get("content-type", "")
            if "application/json" in content_type:
                error_data = response.json()
                raise HTTPException(status_code=500, detail=error_data)
            
            return StreamingResponse(
                response.aiter_bytes(),
                media_type=response.headers.get("content-type", "application/vnd.android.package-archive"),
                headers={
                    "Content-Disposition": f'attachment; filename="reinforced_{taskId}.apk"'
                }
            )
        except Exception as e:
            logger.error(f"下载加固后的APK失败: {str(e)}")
            raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
    